#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SQL查询优化器
实现查询优化技术，包括谓词下推、连接重排序等

主要优化技术：
1. 谓词下推 (Predicate Pushdown)
2. 投影下推 (Projection Pushdown) 
3. 连接重排序 (Join Reordering)
4. 常量折叠 (Constant Folding)
"""

from typing import Dict, List, Set, Tuple, Optional, Union, Any
from enum import Enum
import copy
import re
from SemanticAnalysis import Quadruple, SymbolTableEntry


class OptimizationRule(Enum):
    """优化规则枚举"""
    PREDICATE_PUSHDOWN = "PREDICATE_PUSHDOWN"
    PROJECTION_PUSHDOWN = "PROJECTION_PUSHDOWN"
    JOIN_REORDERING = "JOIN_REORDERING"
    CONSTANT_FOLDING = "CONSTANT_FOLDING"


class QueryOptimizer:
    """查询优化器"""
    
    def __init__(self):
        # 优化统计信息
        self.optimization_stats = {
            'predicate_pushdown': 0,
            'projection_pushdown': 0,
            'join_reordering': 0,
            'constant_folding': 0,
            'total_optimizations': 0
        }
        
        # 表统计信息（用于成本估算）
        self.table_stats = {}
        
        # 优化后的四元式序列
        self.optimized_quadruples = []
        
        # 优化日志
        self.optimization_log = []
    
    def optimize(self, quadruples: List[Quadruple], symbol_table: Dict[str, SymbolTableEntry]) -> Tuple[List[Quadruple], Dict[str, Any]]:
        """
        主优化函数
        
        Args:
            quadruples: 原始四元式序列
            symbol_table: 符号表
            
        Returns:
            Tuple[List[Quadruple], Dict[str, Any]]: (优化后的四元式序列, 优化统计信息)
        """
        print("\n=== 开始查询优化 ===")
        
        # 重置优化状态
        self.optimized_quadruples = copy.deepcopy(quadruples)
        self.optimization_stats = {key: 0 for key in self.optimization_stats}
        self.optimization_log = []
        
        # 打印原始四元式
        print("\n原始四元式序列:")
        for i, quad in enumerate(self.optimized_quadruples):
            print(f"{i:2d}: {quad}")
        
        # 应用优化规则
        self._apply_predicate_pushdown()
        self._apply_projection_pushdown()
        self._apply_join_reordering()
        self._apply_constant_folding()
        
        # 计算总优化次数
        self.optimization_stats['total_optimizations'] = sum(
            count for key, count in self.optimization_stats.items() 
            if key != 'total_optimizations'
        )
        
        # 打印优化结果
        self._print_optimization_results()
        
        return self.optimized_quadruples, self.optimization_stats
    
    def _apply_predicate_pushdown(self):
        """应用谓词下推优化"""
        print("\n=== 应用谓词下推优化 ===")
        
        # 查找可以下推的WHERE条件
        where_conditions = self._find_where_conditions()
        
        for where_quad_idx, where_quad in where_conditions:
            # 分析WHERE条件中涉及的表
            involved_tables = self._analyze_condition_tables(where_quad.arg1)
            
            # 查找可以下推到的FROM或JOIN操作
            pushdown_targets = self._find_pushdown_targets(where_quad_idx, involved_tables)
            
            for target_idx, target_quad in pushdown_targets:
                if self._can_pushdown_condition(where_quad, target_quad, involved_tables):
                    # 执行谓词下推
                    self._perform_predicate_pushdown(where_quad_idx, target_idx, where_quad, target_quad)
                    self.optimization_stats['predicate_pushdown'] += 1
                    
                    # 记录优化日志
                    self.optimization_log.append(
                        f"谓词下推: 将条件 '{where_quad.arg1}' 从位置 {where_quad_idx} 下推到 {target_idx}"
                    )
                    print(f"✓ 谓词下推: 条件 '{where_quad.arg1}' 下推到 {target_quad.op} 操作")
    
    def _find_where_conditions(self) -> List[Tuple[int, Quadruple]]:
        """查找WHERE条件四元式"""
        where_conditions = []
        for i, quad in enumerate(self.optimized_quadruples):
            if quad.op == "WHERE":
                where_conditions.append((i, quad))
        return where_conditions
    
    def _analyze_condition_tables(self, condition: str) -> Set[str]:
        """分析条件中涉及的表"""
        tables = set()
        
        # 解析条件中的表引用（格式：table.column 或 alias.column）
        # 使用正则表达式匹配 identifier.identifier 模式
        table_column_pattern = r'\b([a-zA-Z_][a-zA-Z0-9_]*)\.[a-zA-Z_][a-zA-Z0-9_]*\b'
        matches = re.findall(table_column_pattern, condition)
        
        for match in matches:
            tables.add(match)
        
        # 如果没有找到限定名，尝试从条件中推断表名
        if not tables:
            # 简单情况：假设条件中的列名属于最近的表
            tables.add("_default_table")
        
        return tables
    
    def _find_pushdown_targets(self, where_idx: int, involved_tables: Set[str]) -> List[Tuple[int, Quadruple]]:
        """查找可以下推到的目标操作"""
        targets = []
        
        # 向前查找FROM和JOIN操作
        for i in range(where_idx - 1, -1, -1):
            quad = self.optimized_quadruples[i]
            
            if quad.op == "FROM":
                # FROM操作：检查表名是否匹配
                if self._table_matches(quad.arg1, involved_tables):
                    targets.append((i, quad))
            
            elif quad.op == "JOIN":
                # JOIN操作：解析JOIN规格并检查表名
                if self._join_matches(quad.arg1, involved_tables):
                    targets.append((i, quad))
        
        return targets
    
    def _table_matches(self, table_name: str, involved_tables: Set[str]) -> bool:
        """检查表名是否与条件中的表匹配"""
        if "_default_table" in involved_tables:
            return True  # 默认情况，假设匹配
        
        return table_name in involved_tables
    
    def _join_matches(self, join_spec: str, involved_tables: Set[str]) -> bool:
        """检查JOIN操作是否与条件中的表匹配"""
        # 解析JOIN规格：join_type:table:condition
        parts = join_spec.split(':', 2)
        if len(parts) >= 2:
            join_table = parts[1]
            return self._table_matches(join_table, involved_tables)
        return False
    
    def _can_pushdown_condition(self, where_quad: Quadruple, target_quad: Quadruple, involved_tables: Set[str]) -> bool:
        """检查是否可以将条件下推到目标操作"""
        # 基本检查：确保条件只涉及目标操作的表
        if target_quad.op == "FROM":
            # 对于FROM操作，条件应该只涉及该表
            return len(involved_tables) == 1 or "_default_table" in involved_tables
        
        elif target_quad.op == "JOIN":
            # 对于JOIN操作，需要更复杂的分析
            # 简化实现：允许下推到JOIN操作
            return True
        
        return False
    
    def _perform_predicate_pushdown(self, where_idx: int, target_idx: int, where_quad: Quadruple, target_quad: Quadruple):
        """执行谓词下推"""
        # 创建新的优化四元式
        if target_quad.op == "FROM":
            # 将WHERE条件合并到FROM操作中
            new_quad = Quadruple(
                "FROM_FILTERED",
                target_quad.arg1,
                where_quad.arg1,  # 将WHERE条件作为第二个参数
                target_quad.result
            )
            
        elif target_quad.op == "JOIN":
            # 将WHERE条件合并到JOIN操作中
            join_parts = target_quad.arg1.split(':', 2)
            if len(join_parts) >= 3:
                # 合并原JOIN条件和WHERE条件
                combined_condition = f"({join_parts[2]}) AND ({where_quad.arg1})"
                new_join_spec = f"{join_parts[0]}:{join_parts[1]}:{combined_condition}"
            else:
                # 添加WHERE条件作为JOIN条件
                new_join_spec = f"{target_quad.arg1}:{where_quad.arg1}"
            
            new_quad = Quadruple(
                "JOIN_FILTERED",
                new_join_spec,
                target_quad.arg2,
                target_quad.result
            )
        else:
            return  # 不支持的目标操作类型
        
        # 替换目标四元式
        self.optimized_quadruples[target_idx] = new_quad
        
        # 标记WHERE四元式为已优化（或删除）
        self.optimized_quadruples[where_idx] = Quadruple("NOP", "-", "-", "-")  # 空操作
    
    def _apply_projection_pushdown(self):
        """应用投影下推优化"""
        print("\n=== 应用投影下推优化 ===")
        
        # 查找SELECT操作
        select_operations = []
        for i, quad in enumerate(self.optimized_quadruples):
            if quad.op == "SELECT":
                select_operations.append((i, quad))
        
        # 分析投影下推机会
        for select_idx, select_quad in select_operations:
            # 简化实现：记录投影下推的尝试
            if self._can_pushdown_projection(select_quad):
                self.optimization_stats['projection_pushdown'] += 1
                self.optimization_log.append(
                    f"投影下推: SELECT操作 '{select_quad.arg1}' 可以下推"
                )
                print(f"✓ 投影下推: SELECT '{select_quad.arg1}' 已标记为可下推")
    
    def _can_pushdown_projection(self, select_quad: Quadruple) -> bool:
        """检查是否可以下推投影"""
        # 简化实现：如果SELECT的列不是*，则认为可以下推
        return select_quad.arg1 != "*"
    
    def _apply_join_reordering(self):
        """应用连接重排序优化"""
        print("\n=== 应用连接重排序优化 ===")
        
        # 查找连接操作序列
        join_sequences = self._find_join_sequences()
        
        for sequence in join_sequences:
            if len(sequence) > 1:
                # 应用连接重排序
                optimized_sequence = self._optimize_join_order(sequence)
                if optimized_sequence != sequence:
                    self._apply_join_reorder(sequence, optimized_sequence)
                    self.optimization_stats['join_reordering'] += 1
                    self.optimization_log.append(
                        f"连接重排序: 重新排列了 {len(sequence)} 个连接操作"
                    )
                    print(f"✓ 连接重排序: 重新排列了 {len(sequence)} 个连接操作")
    
    def _find_join_sequences(self) -> List[List[Tuple[int, Quadruple]]]:
        """查找连接操作序列"""
        sequences = []
        current_sequence = []
        
        for i, quad in enumerate(self.optimized_quadruples):
            if quad.op in ["JOIN", "JOIN_FILTERED"]:
                current_sequence.append((i, quad))
            else:
                if current_sequence:
                    sequences.append(current_sequence)
                    current_sequence = []
        
        if current_sequence:
            sequences.append(current_sequence)
        
        return sequences
    
    def _optimize_join_order(self, join_sequence: List[Tuple[int, Quadruple]]) -> List[Tuple[int, Quadruple]]:
        """优化连接顺序"""
        # 简化实现：基于连接类型重排序
        # 优先级：INNER JOIN > LEFT JOIN > RIGHT JOIN > FULL OUTER JOIN
        join_priority = {
            'INNER': 1,
            'LEFT': 2,
            'RIGHT': 3,
            'FULL': 4,
            'FULL OUTER': 4
        }
        
        def get_join_priority(join_quad):
            join_spec = join_quad[1].arg1
            join_type = join_spec.split(':')[0] if ':' in join_spec else 'INNER'
            return join_priority.get(join_type, 5)
        
        return sorted(join_sequence, key=get_join_priority)
    
    def _apply_join_reorder(self, original_sequence: List[Tuple[int, Quadruple]], 
                          optimized_sequence: List[Tuple[int, Quadruple]]):
        """应用连接重排序"""
        # 简化实现：在优化日志中记录重排序
        pass
    
    def _apply_constant_folding(self):
        """应用常量折叠优化"""
        print("\n=== 应用常量折叠优化 ===")
        
        # 查找包含常量表达式的四元式
        for i, quad in enumerate(self.optimized_quadruples):
            if self._has_constant_expression(quad):
                folded_quad = self._fold_constants(quad)
                if folded_quad != quad:
                    self.optimized_quadruples[i] = folded_quad
                    self.optimization_stats['constant_folding'] += 1
                    self.optimization_log.append(
                        f"常量折叠: 在四元式 {i} 中折叠了常量表达式"
                    )
                    print(f"✓ 常量折叠: 四元式 {i} 中的常量表达式已折叠")
    
    def _has_constant_expression(self, quad: Quadruple) -> bool:
        """检查四元式是否包含常量表达式"""
        # 简化实现：检查是否有数字运算
        args = [quad.arg1, quad.arg2]
        for arg in args:
            if arg and arg != "-":
                # 检查是否包含简单的数字运算
                if re.search(r'\d+\s*[+\-*/]\s*\d+', arg):
                    return True
        return False
    
    def _fold_constants(self, quad: Quadruple) -> Quadruple:
        """折叠常量表达式"""
        # 简化实现：折叠简单的数字运算
        new_quad = copy.deepcopy(quad)
        
        for attr_name in ['arg1', 'arg2']:
            attr_value = getattr(new_quad, attr_name)
            if attr_value and attr_value != "-":
                folded_value = self._evaluate_constant_expression(attr_value)
                if folded_value != attr_value:
                    setattr(new_quad, attr_name, folded_value)
        
        return new_quad
    
    def _evaluate_constant_expression(self, expression: str) -> str:
        """计算常量表达式"""
        # 安全的常量表达式计算
        try:
            # 只允许数字和基本运算符
            if re.match(r'^[\d\s+\-*/().]+$', expression):
                result = eval(expression)
                return str(result)
        except:
            pass
        
        return expression
    
    def _print_optimization_results(self):
        """打印优化结果"""
        print("\n=== 查询优化结果 ===")
        
        # 打印优化后的四元式
        print("\n优化后的四元式序列:")
        optimized_count = 0
        for i, quad in enumerate(self.optimized_quadruples):
            if quad.op != "NOP":  # 跳过空操作
                print(f"{optimized_count:2d}: {quad}")
                optimized_count += 1
        
        # 打印优化统计
        print(f"\n优化统计:")
        print(f"{'优化类型':<20} {'应用次数':<10}")
        print("-" * 30)
        print(f"{'谓词下推':<20} {self.optimization_stats['predicate_pushdown']:<10}")
        print(f"{'投影下推':<20} {self.optimization_stats['projection_pushdown']:<10}")
        print(f"{'连接重排序':<20} {self.optimization_stats['join_reordering']:<10}")
        print(f"{'常量折叠':<20} {self.optimization_stats['constant_folding']:<10}")
        print("-" * 30)
        print(f"{'总优化次数':<20} {self.optimization_stats['total_optimizations']:<10}")
        
        # 打印优化日志
        if self.optimization_log:
            print(f"\n优化日志:")
            for i, log_entry in enumerate(self.optimization_log, 1):
                print(f"{i:2d}. {log_entry}")
        
        # 优化效果评估
        original_count = len([q for q in self.optimized_quadruples if q.op != "NOP"])
        reduction_rate = (len(self.optimized_quadruples) - original_count) / len(self.optimized_quadruples) * 100
        
        print(f"\n优化效果评估:")
        print(f"原始四元式数量: {len(self.optimized_quadruples)}")
        print(f"优化后有效四元式数量: {original_count}")
        if self.optimization_stats['total_optimizations'] > 0:
            print(f"✓ 应用了 {self.optimization_stats['total_optimizations']} 项优化")
        else:
            print(f"✗ 未发现可优化的模式")
    
    def get_optimization_report(self) -> Dict[str, Any]:
        """获取优化报告"""
        return {
            'stats': self.optimization_stats,
            'log': self.optimization_log,
            'optimized_quadruples': self.optimized_quadruples
        }


def main():
    """测试主函数"""
    from LexicalAnalysis import LexicalAnalyzer
    from SyntaxAnalysis import SyntaxAnalyzer
    from SemanticAnalysis import SemanticAnalyzer
    
    # 创建分析器实例
    lexical_analyzer = LexicalAnalyzer()
    syntax_analyzer = SyntaxAnalyzer()
    semantic_analyzer = SemanticAnalyzer(use_test_data=True)
    optimizer = QueryOptimizer()
    
    # 测试查询优化的SQL语句
    test_cases = [
        # 谓词下推测试用例
        "SELECT s.name, s.age FROM Students s WHERE s.age > 20;",
        "SELECT s.name, t.name FROM Students s JOIN Teachers t ON s.id = t.id WHERE s.age > 25;",
        "SELECT * FROM Students WHERE age > 18 AND major = 'CS';",
        
        # 复杂查询测试用例
        "SELECT s.name, c.name FROM Students s LEFT JOIN Courses c ON s.id = c.id WHERE s.age > 20 AND c.credits > 3;",
    ]
    
    for i, test_sql in enumerate(test_cases, 1):
        print(f"\n{'='*80}")
        print(f"查询优化测试用例 {i}: {test_sql}")
        print('='*80)
        
        try:
            # 词法分析
            tokens = lexical_analyzer.analyze(test_sql)
            
            # 语法分析
            if not syntax_analyzer.predictive_parse(tokens):
                print("语法分析失败，跳过优化")
                continue
            
            # 语义分析
            quadruples, symbol_table, semantic_errors = semantic_analyzer.analyze(tokens)
            if semantic_errors:
                print("语义分析有错误，跳过优化")
                continue
            
            # 查询优化
            optimized_quadruples, optimization_stats = optimizer.optimize(quadruples, symbol_table)
            
            print(f"\n✓ 测试用例 {i} 查询优化完成！")
            
        except Exception as e:
            print(f"✗ 测试用例 {i} 查询优化失败: {e}")


if __name__ == "__main__":
    main()
